#!/usr/bin/env python3

# 
# Run the test, compare results against the benchmark
#

from __future__ import print_function
try:
  from builtins import str
except:
  pass

# Variables to compare
vars = ['flag0', 'flag3', 'flagis', 'flagos',
        'flag0a', 'flag3a', 'flagisa', 'flagosa',
        'flag0ac', 'flag3ac','flagisac', 'flagosac',
        'flag0ad', 'flag3ad', 'flagisad', 'flagosad']  
tol = 1e-6                  # Absolute tolerance

from boututils.run_wrapper import shell, shell_safe, launch_safe
from boutdata.collect import collect
import numpy as np
from sys import stdout, exit



print("Making Laplacian inversion test")
shell("rm test_laplace")
shell_safe("make > make.log")

# Read benchmark values
print("Reading benchmark data")
bmk = {}
for v in vars:
  bmk[v] = collect(v, path="data", prefix="benchmark", info=False)

print("Running Laplacian inversion test")
success = True

for nproc in [1,2,4]:
  nxpe = 1
  if nproc > 2:
    nxpe = 2
  
  cmd = "./test_laplace nxpe=" + str(nxpe)
  
  shell("rm data/BOUT.dmp.*.nc")

  print("   %d processors (nxpe = %d)...." % (nproc, nxpe))
  s, out = launch_safe(cmd, nproc=nproc, mthread=1, pipe=True)
  with open("run.log."+str(nproc), "w") as f:
    f.write(out)

   # Collect output data
  for v in vars:
    stdout.write("      Checking variable "+v+" ... ")
    result = collect(v, path="data", info=False)
    # Compare benchmark and output
    if np.shape(bmk[v]) != np.shape(result):
      print("Fail, wrong shape")
      success = False
    diff =  np.max(np.abs(bmk[v] - result))
    if diff > tol:
      print("Fail, maximum difference = "+str(diff))
      success = False
    else:
      print("Pass")

if success:
  print(" => All Laplacian inversion tests passed")
  exit(0)
else:
  print(" => Some failed tests")
  exit(1)
